from utils.handian_data import HandianData
from word_vec_encoders.mini_sequence_conv_encoder import MiniSequenceConvEncoder
import pypinyin
import string


class PinyinWordVecEncoder(MiniSequenceConvEncoder):
    """
    该编码器将字符转化为拼音串，补充到固定长度并用卷积和池化提取信息作为字符向量。
    """

    def __init__(self, embedding_dim:int = 32) -> None:        
        # 拼音字符表，包括a-z的字母，1-4的数字，以及末尾填充用的空格
        symbols = [' '] + [c for c in string.ascii_lowercase] + [str(i) for i in range(1, 5)]
        super().__init__(name = 'pinyin', miniseq_symbols = symbols, fixed_miniseq_len = 7, embedding_dim = embedding_dim)

    # def sentence_to_miniseqs_cached(self, sentence:str) -> list:
    #     if sentence in self._sentence_miniseq_cache:
    #         return self._sentence_miniseq_cache[sentence]
    #     sentence_pinyin = pypinyin.pinyin(sentence, style=pypinyin.Style.FINALS_TONE3, errors=lambda x: [[''] for _ in x])
    #     sentence_pinyin = [pinyin[0] for pinyin in sentence_pinyin] # 消除最内层的列表
    #     self._sentence_miniseq_cache[sentence] = sentence_pinyin
    #     return sentence_pinyin

    def sentence_to_miniseqs_cached(self, sentence:str) -> list:
        if sentence in self._sentence_miniseq_cache:
            return self._sentence_miniseq_cache[sentence]
        sentence_miniseqs = []
        for char in sentence:
            miniseq = HandianData.get_pinyin_seq_cached(char)
            miniseq = miniseq.lower()
            sentence_miniseqs.append(miniseq)
        self._sentence_miniseq_cache[sentence] = sentence_miniseqs
        return sentence_miniseqs